package com.aliyun.oss;

// DHKeyExchange.java
import java.math.BigInteger;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

public class DHKeyExchange {

    // 生成 DH 密钥对
    public static KeyPair generateKeyPair() throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("DH");
        SecureRandom random = SecureRandom.getInstanceStrong();
        keyGen.initialize(2048, random); // 使用 2048 位密钥长度
        return keyGen.generateKeyPair();
    }

    // 生成共享密钥
    public static byte[] generateSharedSecret(KeyPair myKeyPair, PublicKey theirPublicKey) throws Exception {
        KeyAgreement keyAgreement = KeyAgreement.getInstance("DH");
        keyAgreement.init(myKeyPair.getPrivate());
        keyAgreement.doPhase(theirPublicKey, true);
        return keyAgreement.generateSecret();
    }

    public static void main(String[] args) {
        try {
            // Alice 生成密钥对
            KeyPair aliceKeyPair = generateKeyPair();
            PublicKey alicePublicKey = aliceKeyPair.getPublic();
            PrivateKey alicePrivateKey = aliceKeyPair.getPrivate();

            // Bob 生成密钥对
            KeyPair bobKeyPair = generateKeyPair();
            PublicKey bobPublicKey = bobKeyPair.getPublic();
            PrivateKey bobPrivateKey = bobKeyPair.getPrivate();

            // Alice 使用 Bob 的公钥生成共享密钥
            byte[] aliceSharedSecret = generateSharedSecret(aliceKeyPair, bobPublicKey);

            // Bob 使用 Alice 的公钥生成共享密钥
            byte[] bobSharedSecret = generateSharedSecret(bobKeyPair, alicePublicKey);

            // 验证共享密钥是否相同
            if (java.util.Arrays.equals(aliceSharedSecret, bobSharedSecret)) {
                System.out.println("密钥协商成功！共享密钥相同。");
                // 打印共享密钥（为了简洁，这里不打印完整的密钥，只打印前16个字节作为示例）
                System.out.println("共享密钥前16个字节: " + java.util.Arrays.toString(java.util.Arrays.copyOfRange(aliceSharedSecret, 0, 16)));
            } else {
                System.out.println("密钥协商失败！共享密钥不同。");
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}